#include "dexmat.h"
#include "dexvec.h"
#include <omp.h>
#include <chrono>
#include <iostream>
#include <algorithm>
#include <iomanip>
#include <fstream>
#include <cstdint>
using namespace std;

void DexMat::init(int rows) {
    vecs.reserve(rows);
    for (int i = 0; i < rows; i++) {
        vecs.push_back(DexVec{});
    }
}

DexVec& DexMat::operator[](int i) {
    return vecs[i];
}

const DexVec& DexMat::operator[](int i) const {
    return vecs[i];
}

size_t DexMat::size() const {
    return vecs.size();
}

void DexMat::addVec(const std::vector<Lit>& vec) {
    vecs.push_back(vec);
}

bool DexMat::VecComp::operator()(const DexVec& a, const DexVec& b) const {
    size_t minLen = min(a.size(), b.size());
    for (size_t dex = 0; dex < minLen; dex++) {
        if (a[dex] == b[dex])
            continue;
        else if (abs(a[dex]) == abs(b[dex])) {
            return a[dex] < b[dex];
        } else {
            return abs(a[dex]) < abs(b[dex]);
        }
    }
    return a.size() < b.size();
}


void DexMat::sortVecs() {
    auto start = chrono::high_resolution_clock::now();
    
    size_t numVecs = vecs.size();
    size_t chunkSize = max(size_t(1), numVecs / omp_get_max_threads());
    VecComp comp;
    
    #pragma omp parallel
    {
        int threadId = omp_get_thread_num();
        size_t start = threadId * chunkSize;
        size_t end = (threadId == omp_get_num_threads() - 1) ? numVecs : start + chunkSize;
        
        if (start < numVecs) {
            sort(vecs.begin() + start, vecs.begin() + end, comp);
        }
    }
    
    auto afterParallelSort = std::chrono::high_resolution_clock::now();
    auto parallelSortTime = chrono::duration_cast<chrono::milliseconds>(
        afterParallelSort - start).count();
    cout << "Parallel sort time: " << parallelSortTime / 1000.0 << " seconds" << endl;
    
    // Then merge sorted chunks
    vector<DexVec> temp(numVecs);
    for (size_t size = chunkSize; size < numVecs; size *= 2) {
        #pragma omp parallel for schedule(dynamic)
        for (size_t head = 0; head < numVecs; head += 2 * size) {
            size_t body = std::min(head + size, numVecs);
            size_t tail = std::min(head + 2 * size, numVecs);
            merge(
                vecs.begin() + head, vecs.begin() + body,
                vecs.begin() + body, vecs.begin() + tail,
                temp.begin() + head,
                comp
            );
        }
        
        // Copy back
        #pragma omp parallel for schedule(static)
        for (size_t i = 0; i < numVecs; i++) {
            vecs[i] = move(temp[i]);
        }
    }
    
    auto end = chrono::high_resolution_clock::now();
    auto totalTime = chrono::duration_cast<chrono::milliseconds>(
        end - start).count();
    auto mergeTime = chrono::duration_cast<chrono::milliseconds>(
        end - afterParallelSort).count();
    
    cout << "Merge time: " << mergeTime / 1000.0 << " seconds" << endl;
    cout << "Total sorting time: " << totalTime / 1000.0 << " seconds" << endl;
}

void DexMat::resize(size_t size) {
    vecs.resize(size);
}

void saveDexMat(const std::string& filename, const DexMat& mat) {
    std::ofstream out(filename, std::ios::binary);
    if (!out) throw std::runtime_error("Cannot open file for writing: " + filename);
    uint64_t numClauses = mat.size();
    out.write(reinterpret_cast<const char*>(&numClauses), sizeof(numClauses));
    for (size_t i = 0; i < numClauses; ++i) {
        const auto& clause = mat[i];
        uint32_t numLits = clause.size();
        out.write(reinterpret_cast<const char*>(&numLits), sizeof(numLits));
        for (size_t j = 0; j < numLits; ++j) {
            int32_t lit = clause[j];
            out.write(reinterpret_cast<const char*>(&lit), sizeof(lit));
        }
    }
}

void loadDexMat(const std::string& filename, DexMat& mat) {
    std::ifstream in(filename, std::ios::binary);
    if (!in) throw std::runtime_error("Cannot open file for reading: " + filename);
    uint64_t numClauses;
    in.read(reinterpret_cast<char*>(&numClauses), sizeof(numClauses));
    mat.resize(numClauses);
    for (uint64_t i = 0; i < numClauses; ++i) {
        uint32_t numLits;
        in.read(reinterpret_cast<char*>(&numLits), sizeof(numLits));
        std::vector<Lit> lits(numLits);
        for (uint32_t j = 0; j < numLits; ++j) {
            int32_t lit;
            in.read(reinterpret_cast<char*>(&lit), sizeof(lit));
            lits[j] = lit;
        }
        mat[i] = lits;
    }
}